#include "armor_tracker/tracker.hpp"

#include <angles/angles.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/convert.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

#include <rclcpp/logger.hpp>
#include <cfloat>
#include <memory>

namespace rm_aim
{
ArmorTracker::ArmorTracker(double max_match_dist, double max_match_yaw)
: tracked_armor_id(""),
  target_state(Eigen::VectorXd::Zero(9)),
  measurement(Eigen::VectorXd::Zero(4)),
  info_position_diff(0.0),
  info_yaw_diff(0.0),
  tracked_armor_count(ArmorCount::kNormal4),
  dz(0.0),
  another_r(0.26),
  tracking_thres(3),
  lost_thres(3),
  tracking_state_(TrackingState::kLost),
  last_update_stamp_(rclcpp::Clock().now()),
  sampling_time_(0.01),
  max_match_distance_(max_match_dist),
  max_match_yaw_diff_(max_match_yaw),
  detect_count_(0),
  lost_count_(0),
  last_yaw_(0.0),
  sigma2_q_xyz_(20.0),
  sigma2_q_yaw_(100.0),
  sigma2_q_r_(800.0),
  r_xyz_factor_(0.05),
  r_yaw_(0.02)
{
}

void ArmorTracker::initialize(const aim_interfaces::msg::Armors::SharedPtr & armors_msg)
{
  if (armors_msg->armors.empty()) {
    RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "No armors to initialize tracker!");
    return;
  }

  // 选择离图像中心最近的装甲板作为初始目标
  double min_center_dist = DBL_MAX;
  tracked_armor = armors_msg->armors[0];
  for (const auto & armor : armors_msg->armors) {
    if (armor.distance_to_image_center < min_center_dist) {
      min_center_dist = armor.distance_to_image_center;
      tracked_armor = armor;
    }
  }

  // 初始化EKF和追踪参数
  initEKF(tracked_armor);
  tracked_armor_id = tracked_armor.number;
  tracking_state_ = TrackingState::kDetecting;
  updateArmorCount(tracked_armor);

  RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "Tracker initialized with armor ID: %s",
               tracked_armor_id.c_str());
}

void ArmorTracker::updateTracking(const aim_interfaces::msg::Armors::SharedPtr & armors_msg)
{
  // 计算采样时间（当前时间 - 上一次更新时间）
  rclcpp::Time current_stamp = rclcpp::Clock().now();
  sampling_time_ = (current_stamp - last_update_stamp_).seconds();
  last_update_stamp_ = current_stamp;
  // 限制采样时间范围（避免异常值）
  if (sampling_time_ <= 0 || sampling_time_ > 0.1) {
    sampling_time_ = 0.01;
  }

  // EKF预测：构造状态转移矩阵F和过程噪声Q
  Eigen::MatrixXd F(9, 9);  // 状态转移矩阵（9x9）
  F.setIdentity();
  F(0, 1) = sampling_time_;  // xc = xc_prev + vxc*dt
  F(2, 3) = sampling_time_;  // yc = yc_prev + vyc*dt
  F(4, 5) = sampling_time_;  // za = za_prev + vza*dt
  F(6, 7) = sampling_time_;  // yaw = yaw_prev + vyaw*dt

  Eigen::MatrixXd Q(9, 9);  // 过程噪声矩阵（9x9）
  Q.setZero();
  double t = sampling_time_;
  // 位置/速度的过程噪声（基于sigma2_q_xyz_）
  double q_pos = pow(t, 4) / 4 * sigma2_q_xyz_;
  double q_pos_vel = pow(t, 3) / 2 * sigma2_q_xyz_;
  double q_vel = pow(t, 2) * sigma2_q_xyz_;
  // 偏航角/角速度的过程噪声（基于sigma2_q_yaw_）
  double q_yaw = pow(t, 4) / 4 * sigma2_q_yaw_;
  double q_yaw_vyaw = pow(t, 3) / 2 * sigma2_q_yaw_;
  double q_vyaw = pow(t, 2) * sigma2_q_yaw_;
  // 半径的过程噪声（基于sigma2_q_r_）
  double q_r = pow(t, 4) / 4 * sigma2_q_r_;

  // 填充过程噪声矩阵Q
  Q << q_pos,    q_pos_vel, 0,        0,        0,        0,        0,        0,        0,
       q_pos_vel, q_vel,    0,        0,        0,        0,        0,        0,        0,
       0,         0,        q_pos,    q_pos_vel, 0,        0,        0,        0,        0,
       0,         0,        q_pos_vel, q_vel,    0,        0,        0,        0,        0,
       0,         0,        0,        0,        q_pos,    q_pos_vel, 0,        0,        0,
       0,         0,        0,        0,        q_pos_vel, q_vel,    0,        0,        0,
       0,         0,        0,        0,        0,        0,        q_yaw,    q_yaw_vyaw, 0,
       0,         0,        0,        0,        0,        0,        q_yaw_vyaw, q_vyaw,    0,
       0,         0,        0,        0,        0,        0,        0,        0,        q_r;

  // 执行EKF预测
  Eigen::VectorXd ekf_prediction = ekf_.predict(F, Q);
  target_state = ekf_prediction;  // 初始用预测结果作为目标状态
  RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "EKF predicted: x=%.2f, y=%.2f",
               ekf_prediction(0), ekf_prediction(2));

  bool matched = false;  // 标记是否匹配到装甲板
  info_position_diff = DBL_MAX;
  info_yaw_diff = DBL_MAX;

  if (!armors_msg->armors.empty()) {
    // 寻找同ID装甲板中与预测位置最接近的目标
    aim_interfaces::msg::Armor same_id_armor;
    int same_id_count = 0;
    Eigen::Vector3d predicted_armor_pos = getArmorPosFromState(ekf_prediction);
    double min_pos_diff = DBL_MAX;
    double min_yaw_diff = DBL_MAX;

    for (const auto & armor : armors_msg->armors) {
      if (armor.number == tracked_armor_id) {  // 仅考虑同ID装甲板
        same_id_armor = armor;
        same_id_count++;

        // 计算当前装甲板与预测位置的距离
        Eigen::Vector3d current_pos(
          armor.pose.position.x, armor.pose.position.y, armor.pose.position.z);
        double pos_diff = (predicted_armor_pos - current_pos).norm();

        if (pos_diff < min_pos_diff) {  // 更新最近装甲板
          min_pos_diff = pos_diff;
          double current_yaw = quaternionToYaw(armor.pose.orientation);
          min_yaw_diff = std::abs(current_yaw - ekf_prediction(6));
          tracked_armor = armor;
        }
      }
    }

    // 记录调试信息
    info_position_diff = min_pos_diff;
    info_yaw_diff = min_yaw_diff;

    // 检查是否匹配成功（距离和角度差在阈值内）
    if (min_pos_diff < max_match_distance_ && min_yaw_diff < max_match_yaw_diff_) {
      matched = true;
      // 构造观测值（x,y,z,yaw）
      Eigen::Vector3d measured_pos(
        tracked_armor.pose.position.x, tracked_armor.pose.position.y,
        tracked_armor.pose.position.z);
      double measured_yaw = quaternionToYaw(tracked_armor.pose.orientation);
      measurement = Eigen::Vector4d(measured_pos.x(), measured_pos.y(), measured_pos.z(), measured_yaw);

      // 构造观测雅可比矩阵H和观测噪声R
      Eigen::MatrixXd H(4, 9);  // 观测矩阵（4x9）
      H.setZero();
      double yaw = ekf_prediction(6);
      double r = ekf_prediction(8);
      // 第0行：xa关于状态的偏导
      H(0, 0) = 1.0;
      H(0, 6) = r * std::sin(yaw);
      H(0, 8) = -std::cos(yaw);
      // 第1行：ya关于状态的偏导
      H(1, 2) = 1.0;
      H(1, 6) = -r * std::cos(yaw);
      H(1, 8) = -std::sin(yaw);
      // 第2行：za关于状态的偏导
      H(2, 4) = 1.0;

      // 第3行：yaw关于状态的偏导
      H(3, 6) = 1.0;

      /*
      * 观测噪声矩阵R的解释
      * 当目标距离相机较远时，同样的像素误差会导致更大的位置误差
      * 当目标距离相机较近时，位置估计相对更准确，噪声应该更小
      */
      Eigen::MatrixXd R(4, 4);  // 观测噪声矩阵（4x4）
      R.setZero();
      R(0, 0) = std::abs(r_xyz_factor_ * measurement(0));  // x方向观测噪声
      R(1, 1) = std::abs(r_xyz_factor_ * measurement(1));  // y方向观测噪声
      R(2, 2) = std::abs(r_xyz_factor_ * measurement(2));  // z方向观测噪声
      R(3, 3) = r_yaw_;                                   // 偏航角观测噪声

      // 观测函数：从状态映射到观测值
      auto observation_func = [](const Eigen::VectorXd & x) {
      double xc = x(0), yc = x(2);
      double yaw = x(6), r = x(8);
      double xa = xc - r * std::cos(yaw);
      double ya = yc - r * std::sin(yaw);
      return Eigen::Vector4d(xa, ya, x(4), yaw);
      };

       // 防止夹角求差出现异常值
      auto z_subtract = [](const Eigen::VectorXd & a, const Eigen::VectorXd & b) -> Eigen::VectorXd {
      Eigen::VectorXd c = a - b;
      c[3] = angles::normalize_angle(c[3]);  // 处理yaw残差（观测向量第3位是yaw）
      return c;
      };

      // 执行EKF更新
      target_state = ekf_.update(
      measurement,    // 观测值z
      H,              // 观测雅可比矩阵H
      R,              // 观测噪声矩阵R
      observation_func,  // 非线性观测模型h(x)
      z_subtract
      );

      // 输出chi-square验证结果
      logChiSquareValidation();
    }
    // 处理装甲板跳变（同ID但角度差过大）
    else if (same_id_count == 1 && min_yaw_diff > max_match_yaw_diff_) {
      handleArmorJump(same_id_armor);
    }
    // 未匹配到装甲板
    else {
      RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "No matched armor (ID: %s)",
                  tracked_armor_id.c_str());
    }
  }

  // 限制半径范围（防止EKF估计发散）
  if (target_state(8) < 0.12) {
    target_state(8) = 0.12;
    ekf_.x = target_state;
  } else if (target_state(8) > 0.4) {
    target_state(8) = 0.4;
    ekf_.x = target_state;
  }

  // 约束v_yaw（角速度）波动范围
  if (tracked_armor_count == ArmorCount::kOutpost3) {
    double min_v_yaw = 0.3;  // 前哨站最小旋转角速度（rad/s，实测）
    double max_v_yaw = 0.5;  // 最大旋转角速度
    // 顺时针和逆时针旋转
    if (target_state(7) > 0) {
      target_state(7) = std::clamp(target_state(7), min_v_yaw, max_v_yaw);
    } else {
      target_state(7) = std::clamp(target_state(7), -max_v_yaw, -min_v_yaw);
    }
    ekf_.x(7) = target_state(7);  // 同步到EKF状态
  }

  // 更新追踪状态机
  switch (tracking_state_) {
    case TrackingState::kDetecting:
      if (matched) {
        detect_count_++;
        if (detect_count_ > tracking_thres) {  // 连续匹配次数达标，进入稳定追踪
          tracking_state_ = TrackingState::kTracking;
          detect_count_ = 0;
          RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "Enter tracking state (ID: %s)",
                      tracked_armor_id.c_str());
        }
      } else {  // 未匹配，回到丢失状态
        tracking_state_ = TrackingState::kLost;
        detect_count_ = 0;
      }
      break;

    case TrackingState::kTracking:
      if (!matched) {  // 未匹配，进入暂时丢失
        tracking_state_ = TrackingState::kTempLost;
        lost_count_ = 1;
        RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "Temporarily lost target (ID: %s)",
                    tracked_armor_id.c_str());
      }
      break;

    case TrackingState::kTempLost:
      if (!matched) {  // 持续丢失，计数累加
        lost_count_++;
        if (lost_count_ > lost_thres) {  // 超过丢失阈值，回到丢失状态
          tracking_state_ = TrackingState::kLost;
          lost_count_ = 0;
          RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "Lost target (ID: %s)",
                       tracked_armor_id.c_str());
        }
      } else {  // 重新匹配，回到追踪状态
        tracking_state_ = TrackingState::kTracking;
        lost_count_ = 0;
        RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "Recovered target (ID: %s)",
                    tracked_armor_id.c_str());
      }
      break;

    case TrackingState::kLost:
      // 丢失状态下无需处理，等待重新初始化
      break;
  }
}

void ArmorTracker::setEKFParams(
  double sigma_q_xyz, double sigma_q_yaw, double sigma_q_r,
  double r_xyz_factor, double r_yaw)
{
  sigma2_q_xyz_ = sigma_q_xyz;
  sigma2_q_yaw_ = sigma_q_yaw;
  sigma2_q_r_ = sigma_q_r;
  r_xyz_factor_ = r_xyz_factor;
  r_yaw_ = r_yaw;
}

void ArmorTracker::initEKF(const aim_interfaces::msg::Armor & armor)
{
  // 提取装甲板位置和角度
  double armor_x = armor.pose.position.x;
  double armor_y = armor.pose.position.y;
  double armor_z = armor.pose.position.z;
  last_yaw_ = 0.0;
  double yaw = quaternionToYaw(armor.pose.orientation);

  // 计算目标中心初始位置（假设装甲板到中心距离为0.2m）
  double radius;
  if(armor.number == "outpost")
  {
    radius = 0.2765;
  }else{
    radius = 0.2;
  }

  double center_x = armor_x + radius * std::cos(yaw);  // 中心x坐标
  double center_y = armor_y + radius * std::sin(yaw);  // 中心y坐标

  // 初始化目标状态向量（9维）
  // [xc, vxc, yc, vyc, za, vza, yaw, vyaw, r]
  Eigen::VectorXd init_state(9);
  init_state << center_x, 0.0,    // x坐标及速度
                center_y, 0.0,    // y坐标及速度
                armor_z,  0.0,    // z坐标及速度
                yaw,      0.0,    // 偏航角及角速度
                radius;           // 装甲板到中心距离

  // 偏航角限制的状态求和函数
  auto x_add = [this](const Eigen::VectorXd & a, const Eigen::VectorXd & b) -> Eigen::VectorXd {
    Eigen::VectorXd c = a + b;
    c[6] = angles::normalize_angle(c[6]);   // 限制偏航角在[-π, π]
    return c;
  };

  // 初始化EKF
  ekf_ = tools::ExtendedKalmanFilter(init_state, init_covariance_, x_add);
  target_state = init_state;
  another_r = radius;  // 初始化备用半径
  RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), "EKF initialized with center (%.2f, %.2f)",
               center_x, center_y);
}

void ArmorTracker::handleArmorJump(const aim_interfaces::msg::Armor & current_armor)
{
  double current_yaw = quaternionToYaw(current_armor.pose.orientation);
  target_state(6) = current_yaw;  // 更新偏航角
  updateArmorCount(current_armor);

  // 4装甲板目标需要切换半径和高度
  if (tracked_armor_count == ArmorCount::kNormal4) {
    dz = target_state(4) - current_armor.pose.position.z;  // 记录高度差
    target_state(4) = current_armor.pose.position.z;       // 更新z坐标
    std::swap(target_state(8), another_r);                 // 切换半径
  }

  RCLCPP_WARN(rclcpp::get_logger("armor_tracker"), "Armor jumped (ID: %s), new radius: %.2f",
              tracked_armor_id.c_str(), target_state(8));

  // 若位置偏差过大，重置EKF状态
  Eigen::Vector3d current_pos(
    current_armor.pose.position.x, current_armor.pose.position.y,
    current_armor.pose.position.z);
  Eigen::Vector3d estimated_pos = getArmorPosFromState(target_state);
  if ((current_pos - estimated_pos).norm() > max_match_distance_) {
    double radius = target_state(8);
    // 重置目标中心位置
    target_state(0) = current_pos.x() + radius * std::cos(current_yaw);  // xc
    target_state(2) = current_pos.y() + radius * std::sin(current_yaw);  // yc
    target_state(4) = current_pos.z();                                   // za
    target_state.segment(1, 3) << 0.0, 0.0, 0.0;  // 重置速度为0
    ekf_.x = target_state;
    RCLCPP_ERROR(rclcpp::get_logger("armor_tracker"), "EKF state reset due to large position error");
  } else {
    ekf_.x = target_state;
  }
}

void ArmorTracker::updateArmorCount(const aim_interfaces::msg::Armor & armor)
{
  // 平衡步兵（ID 3/4/5，大型装甲板）→ 2个装甲板
  if (armor.type == "large_armor" && 
      (tracked_armor_id == "3" || tracked_armor_id == "4" || tracked_armor_id == "5")) {
    tracked_armor_count = ArmorCount::kBalance2;
  }
  // 前哨站（ID "outpost"）→ 3个装甲板
  else if (tracked_armor_id == "outpost") {
    tracked_armor_count = ArmorCount::kOutpost3;
  }
  // 其他目标→4个装甲板
  else {
    tracked_armor_count = ArmorCount::kNormal4;
  }
}

double ArmorTracker::quaternionToYaw(const geometry_msgs::msg::Quaternion & q)
{
  // 四元数转RPY角
  tf2::Quaternion tf_quat;
  tf2::fromMsg(q, tf_quat);
  double roll, pitch, yaw;
  tf2::Matrix3x3(tf_quat).getRPY(roll, pitch, yaw);

  // 处理角度连续性（将-pi~pi映射为连续角度）
  yaw = last_yaw_ + angles::shortest_angular_distance(last_yaw_, yaw);
  last_yaw_ = yaw;  // 保存当前角度用于下次计算

  return yaw;
}

Eigen::Vector3d ArmorTracker::getArmorPosFromState(const Eigen::VectorXd & state)
{
  double center_x = state(0);   // 目标中心x
  double center_y = state(2);   // 目标中心y
  double armor_z = state(4);    // 装甲板z坐标
  double yaw = state(6);        // 偏航角
  double radius = state(8);     // 装甲板到中心距离

  // 计算装甲板位置
  double armor_x = center_x - radius * std::cos(yaw);
  double armor_y = center_y - radius * std::sin(yaw);

  return Eigen::Vector3d(armor_x, armor_y, armor_z);
}

void ArmorTracker::logChiSquareValidation() {
    // 记录详细信息
    RCLCPP_DEBUG(rclcpp::get_logger("armor_tracker"), 
      "EKF Chi-Square Validation:\n"
      "  NIS Value: %.6f (Test: %s)\n"
      "  NEES Value: %.6f (Test: %s)\n"
      "  Residuals: [x: %.4f, y: %.4f, z: %.4f, yaw: %.4f]\n"  
      "  Recent Failure Rate: %.2f%%\n"
      "  Total Updates: %d, NIS Failures: %d",
      ekf_.data["nis"], (ekf_.data["nis_fail"] > 0 ? "FAILED" : "PASSED"),
      ekf_.data["nees"], (ekf_.data["nees_fail"] > 0 ? "FAILED" : "PASSED"),
      ekf_.data["residual_x"], ekf_.data["residual_y"],  
      ekf_.data["residual_z"], ekf_.data["residual_yaw"],
      ekf_.data["recent_nis_failures"] * 100,
      ekf_.getTotalCount(), ekf_.getNisCount());
        
    // 如果失败率过高，输出警告
    if (ekf_.data["recent_nis_failures"] > 0.3) {  // 如果近期失败率超过30%
        RCLCPP_WARN(rclcpp::get_logger("armor_tracker"), 
            "High recent NIS failure rate: %.2f%% - potential tracking instability",
            ekf_.data["recent_nis_failures"] * 100);
    }
}

}  // namespace rm_aim